/*
 * Copyright (c) 2025 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Reaching Definitions Data Flow Analysis
 *
 * This module implements the Reaching Definitions data flow analysis algorithm.
 * Reaching Definitions is a forward data flow analysis that determines, for each
 * program point, the set of variable definitions (assignments) that may reach
 * that point without being overwritten.
 *
 * Key Components:
 * 1. **Transfer Function**:
 *    - Computes the out set for each node based on its in set.
 *    - Uses gen and kill sets to model the effects of assignments:
 *      - **gen**: The set of definitions generated by the current node.
 *      - **kill**: The set of definitions killed (overwritten) by the current node.
 *
 * 2. **Meet Operation**:
 *    - Combines data flow values from multiple paths (e.g., union for reaching definitions).
 *    - Ensures that the analysis is conservative (safe) by over-approximating the result.
 *
 *  The analysis is forward, meaning it propagates information from predecessors to successors.
 *
 */

import { ArkAssignStmt, Stmt } from '../base/Stmt';
import { Value } from '../base/Value';
import { BaseImplicitGraph, NodeID } from '../graph/BaseImplicitGraph';
import { ArkMethod } from '../model/ArkMethod';
import { DataFlowProblem, TransferFunction, FlowGraph } from './GenericDataFlow';
import { SparseBitVector } from '../../utils/SparseBitVector';
import { Cfg } from '../graph/Cfg';

type RDNode = Stmt;
type DFNodeCollection = SparseBitVector;
let coCtor: new (s: number) => SparseBitVector = SparseBitVector;
const BV_SIZE = 32;

export class ReachingDefProblem implements DataFlowProblem<NodeID, DFNodeCollection> {
    flowGraph: ReachingDefFlowGraph;
    transferFunction: ReachingDefTransferFunction;
    meet: (a: DFNodeCollection, b: DFNodeCollection) => DFNodeCollection;
    initIn: Map<NodeID, DFNodeCollection>;
    initOut: Map<NodeID, DFNodeCollection>;
    forward: boolean;
    empty: DFNodeCollection = new coCtor(BV_SIZE);

    constructor(method: ArkMethod, forward: boolean = true) {
        this.flowGraph = new ReachingDefFlowGraph(method);
        this.transferFunction = new ReachingDefTransferFunction(this.flowGraph);
        this.meet = (x: DFNodeCollection, y: DFNodeCollection): DFNodeCollection => {
            let r = x.clone();
            r.unionWith(y);
            return r;
        };
        this.initIn = new Map(this.flowGraph.nodesInPostOrder.map(i => [i, new coCtor(BV_SIZE)]));
        this.initOut = new Map(this.flowGraph.nodesInPostOrder.map(i => [i, new coCtor(BV_SIZE)]));
        this.forward = forward;
    }
}

/**
 * Represents the control flow graph (CFG) for reaching definitions analysis.
 * This class implements the FlowGraph interface and provides methods to retrieve
 * successors and predecessors of nodes, as well as topological orderings of nodes.
 */
class ReachingDefFlowGraph extends BaseImplicitGraph<RDNode> implements FlowGraph<NodeID> {
    nodesInPostOrder: NodeID[];

    constructor(method: ArkMethod) {
        super();
        const cfg = method.getCfg();
        if (!cfg) {
            throw new Error('CFG not found');
        }

        const nodes = cfg.getStmts();
        this.nodeToIdMap = new Map(nodes.map((x, i) => [x, i]));
        this.idToNodeMap = new Map(nodes.map((x, i) => [i, x]));
        this.nodesInPostOrder = nodes.map((_, i) => i);

        this.initSuccPred(nodes, cfg);
    }

    getGraphName(): string {
        return 'Reaching Definition Flow Graph';
    }

    dumpNodes(): void {
        this.nodeToIdMap?.forEach((id, node) => console.log(id + ': ' + node.toString()));
    }

    private initSuccPred(nodes: RDNode[], cfg: Cfg): void {
        this.succMap = new Map();
        this.predMap = new Map();

        cfg.getBlocks().forEach(bb => {
            let stmts = bb.getStmts();
            if (stmts.length === 0) {
                return;
            }

            for (let i = 0; i < stmts.length - 1; i++) {
                let c = this.nodeToIdMap!.get(stmts[i])!;
                let n = this.nodeToIdMap!.get(stmts[i + 1])!;
                if (c === undefined || n === undefined) {
                    continue;
                }
                this.succMap.set(c, [n]);
                this.predMap.set(n, [c]);
            }

            let terminate = bb.getTail();
            if (!terminate) {
                throw new Error('cfg has no terminal');
            }

            let successors = bb.getSuccessors();
            // try...catch语句，catch所在的block在CFG表示里是没有前驱block的，需要在这里额外查找并将exceptionalSuccessorBlocks作为try块的后继块之一
            const exceptionalSuccessorBlocks = bb.getExceptionalSuccessorBlocks();
            if (exceptionalSuccessorBlocks !== undefined) {
                successors.push(...exceptionalSuccessorBlocks);
            }
            successors.forEach(succBB => {
                let head = succBB.getHead();
                if (!head) {
                    return;
                }
                let t = this.nodeToIdMap?.get(terminate!)!;
                let h = this.nodeToIdMap?.get(head)!;
                if (t === undefined || h === undefined) {
                    return;
                }
                // Terminate's succ
                let succ = this.succMap.get(t) ?? [];
                succ.push(h);
                this.succMap.set(t, succ);

                // Head's pred
                let pred = this.predMap.get(h) ?? [];
                pred.push(t);
                this.predMap.set(h, pred);
            });
        });
    }
}

/**
 * Represents the transfer function for reaching definitions analysis.
 */
export class ReachingDefTransferFunction implements TransferFunction<NodeID, DFNodeCollection> {
    gen: DFNodeCollection;
    kill: Map<NodeID, DFNodeCollection>;

    constructor(flowGraph: ReachingDefFlowGraph) {
        this.gen = new coCtor(BV_SIZE);
        this.kill = new Map();
        this.initGenKill(flowGraph);
    }

    apply(n: NodeID, x: DFNodeCollection): DFNodeCollection {
        const result = x.clone();
        if (this.gen.test(n)) {
            result.set(n);
        }

        const killSet = this.kill.get(n);
        if (killSet) {
            for (const item of killSet) {
                result.reset(item);
            }
        }
        return result;
    }

    private initGenKill(g: ReachingDefFlowGraph): void {
        let genValue2Nodes: Map<Value, DFNodeCollection> = new Map();
        // Init Gen
        g.getNodeToIdMap().forEach((id, node) => {
            if (node instanceof ArkAssignStmt) {
                let lop = node.getLeftOp();
                let genNodes = genValue2Nodes.get(lop) ?? new coCtor(BV_SIZE);
                genNodes.set(id);
                genValue2Nodes.set(lop, genNodes);
                this.gen.set(id);
            }
        });

        // Init Kill
        genValue2Nodes.forEach((defNodes, v) => {
            for (const i of defNodes) {
                const killSet = defNodes.clone();
                killSet.reset(i);
                this.kill.set(i, killSet);
            }
        });
    }
}
